#%%
%matplotlib inline
import numpy as np, pandas as pd, pathlib, matplotlib.pyplot as plt
grains = ['dn', 'jo', 'mb', 'mk', 'pg', 'sf']
grains_new = {'dn': 'tg', 'jo': 'ng2', 'mb': 'ng3', 'mk': 'ng4', 'pg': 'ng5', 'sf': 'sf'}

# load up settling velocity data and process it

df_density = pd.read_csv('1_var_density_2022.csv').set_index('Unnamed: 0').rename_axis(None)
df_shape = pd.read_csv('2_var_shape_2022.csv').set_index('Unnamed: 0').rename_axis(None)

df_vel = {}
for grain in grains:
    vels = {}
    if (grain == 'dn') or (grain == 'mb') or (grain == 'pg'): date = '20220208'
    else: date = '20220210'

    vels = pd.read_csv('observations/settling_velocity/%s/%s/%s_%i_settling_velocities.csv'%(date, grain, grain, 1))
    for i in range(2,5):
        try:
            temp = pd.read_csv('observations/settling_velocity/%s/%s/%s_%i_settling_velocities.csv'%(date, grain, grain, i))
        except FileNotFoundError: 
            continue
        vels = pd.concat([vels, temp])

    vels = vels.reset_index().drop('index', axis=1).drop('particleId', axis=1)
    df_vel[grain] = {'mean': vels.vel.mean()/1000, 'se': vels.vel.std()/1000/np.sqrt(vels.vel.values.size)}
df_vel = pd.DataFrame(df_vel).T
vel = {grains_new[grain]: df_vel['mean'][grain] for grain in grains}
dvel = {grains_new[grain]: df_vel.se[grain] for grain in grains}

CD = {}
for grain_in in df_vel['mean'].index:
    grain = grains_new[grain_in]
    CD[grain] = 4*(df_density.total_grain_ro[grain]/990 - 1)*9.81*df_shape.dn[grain] / (3*df_vel['mean'][grain_in]**2) * df_shape.CSF[grain]

df = pd.DataFrame({'dvel_mean': dvel, 'vel_mean': vel, 'CD': CD})
df.to_csv('4_var_settling_velocities_2022.csv')
# %%
